#!/usr/bin/env python3

import requests
from tqdm import tqdm
import argparse
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed

def download_file(id):
    try:
        url = f'https://zinc20.docking.org/substances/{id}.smi'
        response = requests.get(url, timeout=10)
        response.raise_for_status()  # 检查请求是否成功
        return response.text
    except requests.RequestException as e:
        print(f"Error downloading {url}: {e}")
        return None

def download_and_merge(urls, output_file, workers):
    with ThreadPoolExecutor(max_workers=workers) as executor:
        # 提交所有任务
        futures = {executor.submit(download_file, url): url for url in urls}
        
        # 使用tqdm显示进度条
        results = []
        for future in tqdm(as_completed(futures), total=len(urls), desc="Downloading"):
            results.append(future.result())
    
    with open(output_file, "w", encoding="utf-8") as f:
        for content in results:
            if content:
                f.write(content)  # 用换行分隔文件内容

def main():
    # Argument parsing
    parser = argparse.ArgumentParser(description="Download SMI file from zinc id list.")
    parser.add_argument("-i", "--input", required=True, help="Input TSV/CSV file with zinc_id column")
    parser.add_argument("-o", "--output", required=True, help="Output smi file")
    parser.add_argument("-w", "--workers", default=10, type=int, help="Maximum number of threads for concurrent execution")
    args = parser.parse_args()

    # Read input TSV file
    try:
        df = pd.read_csv(args.input, sep=None, engine="python")
        if "zinc_id" not in df.columns:
            raise ValueError("Input TSV must contain a 'zinc_id' column.")
        
        # Get unique ZINC IDs
        zinc_ids = df["zinc_id"].unique()
        print(f"Found unique zinc id: {len(zinc_ids)}")
    except Exception as e:
        print(f"Error reading input file: {e}")
        return

    # Query UniChem for each ZINC ID
    download_and_merge(zinc_ids, output_file = args.output, workers = args.workers)

if __name__ == "__main__":
    main()
